import json
import os
import shlex
import subprocess
from typing import Dict, List, Optional, Tuple, TypedDict


class CompileCommandItem(TypedDict):
    """
    TypedDict for a single item in the compile_commands.json file
    """

    command: str
    arguments: List[str]
    directory: str
    file: str


def extract_include_directories(splitted_cmd: List[str]) -> List[str]:
    """
    Extract included directories from a compile command

    :command: compile command to extract included directories from
    :return: list of include directories
    """

    include_directories = []
    for index, cmd_item in enumerate(splitted_cmd):
        if cmd_item == "-I":
            include_directories.append(splitted_cmd[index + 1])
        elif cmd_item.startswith("-I"):
            include_directories.append(splitted_cmd[index].removeprefix("-I"))
    return include_directories


def extract_compile_options(
    splitted_command: List[str], keys: List[str]
) -> Dict[str, str]:
    """
    Extract compile options from a compile command

    :splitted_command: compile command to extract compile options from
    :keys: list of keys to extract
    :return: dictionary of compile options
    """
    splitted = splitted_command
    compile_options = {}
    for index, cmd_item in enumerate(splitted):
        if cmd_item in keys:
            compile_options[cmd_item] = splitted[index + 1]
        elif cmd_item.split("=")[0] in keys:
            compile_options[cmd_item.split("=")[0]] = cmd_item.split("=")[1]
    return compile_options


class Preprocessor:
    """
    Class to preprocess the building data using the compile_commands.json file

    :compile_commands_file: path to compile_commands.json
    :preprocess_path: path to run the preprocessing commands
    """

    def __init__(
        self,
        compile_commands_file: str,
        output_path: str,
        preprocess_path: Optional[str] = None,
    ) -> None:

        self.preprocess_path = (
            preprocess_path
            if preprocess_path is not None
            else os.path.dirname(os.path.abspath(compile_commands_file))
        )
        with open(compile_commands_file) as f:
            self.compile_commands = json.load(f)
        self.output_path = output_path

    def get_compile_of_file(self, filename: str) -> tuple[str, str]:
        # compile_options = {}
        compile_command: CompileCommandItem
        for compile_command in self.compile_commands:
            if os.path.samefile(compile_command["file"], filename):
                _, directories, options = self.extract_command_options(compile_command)
                return compile_command["directory"], self.format_includes_and_options(
                    directories, options
                )
        raise ValueError(f"{filename} not found in compile_commands.json")

    def preprocess(self, mode=".i"):
        assert mode in {".i", ".ll"}
        for compile_command in self.compile_commands:
            filename = os.path.basename(compile_command["file"])
            sourcefile_abs_path = os.path.dirname(
                compile_command["file"]
                if os.path.isabs(compile_command["file"])
                else os.path.join(self.preprocess_path, compile_command["file"])
            )
            # if mode==".i":
            preprocessed_filename = os.path.splitext(filename)[0] + mode
            # Get the output path for preprocessed file
            output_file = os.path.join(
                self.output_path,
                os.path.relpath(sourcefile_abs_path, self.preprocess_path),
                preprocessed_filename,
            )
            if compile_command["directory"] == self.preprocess_path:
                self.preprocess_single_file(compile_command, output_file)

    def preprocess_single_file(
        self, compile_command: CompileCommandItem, output_file: str
    ):
        cmd = self.format_preprocess_single_item_cmd(compile_command, output_file)
        subprocess.run(cmd, shell=True, check=True, cwd=self.preprocess_path)

    def format_includes_and_options(
        self, includes: list[str], options: dict[str, str]
    ) -> str:
        options_str = " ".join([f"{k}={v}" for k, v in options.items()])
        directories_str = " ".join([f"-I{d}" for d in includes])
        return f" {options_str} {directories_str} "

    def format_preprocess_single_item_cmd(
        self, compile_command_item: CompileCommandItem, output_file: str
    ) -> str:
        """
        Format cmd to compile the file using the compile command and return the preprocessed file

        :compile_command_item: compile command item
        :return: preprocessing command
        """
        compiler, directories, options = self.extract_command_options(
            compile_command_item
        )
        output_file_dir = os.path.dirname(output_file)
        if not os.path.exists(output_file_dir):
            os.makedirs(output_file_dir)
        if output_file.endswith(".i"):
            return f"{compiler} -o {output_file} -E {self.format_includes_and_options(directories, options)} {compile_command_item['file']} "
        elif output_file.endswith(".ll"):
            return f"clang -o {output_file} -S -DHAVE_CONFIG_H -emit-llvm {self.format_includes_and_options(directories, options)} {compile_command_item['file']} "
        else:
            raise ValueError(f"Unknown output file extension: {output_file}")

    @classmethod
    def extract_command_options(
        cls, compile_command_item: CompileCommandItem
    ) -> Tuple[str, list[str], dict[str, str]]:
        """
        Extract the include options from the compile command

        :compile_command: dict containing the compile command
        :return: list of include options
        """
        if "arguments" in compile_command_item:
            compiler = compile_command_item["arguments"][0]
            splitted_cmd = compile_command_item["arguments"]
        else:
            compiler = compile_command_item["command"].split(" ", maxsplit=1)[0]
            splitted_cmd = shlex.split(compile_command_item["command"])

        directories = extract_include_directories(splitted_cmd)
        options = extract_compile_options(splitted_cmd, ["-std"])
        return compiler, directories, options
